package org.zhang.core.asm;

import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.AdviceAdapter;
import org.zhang.core.track.*;
import org.zhang.core.util.JacksonUtil;

import java.util.ArrayList;
import java.util.List;

public class MethodInterceptorMethodVisitor extends AdviceAdapter {
    private static final String SYSTEM_CLASS = "java/lang/System";
    private static final String CURRENT_TIME_MILLIS_METHOD = "currentTimeMillis";
    private int startTimeVarIndex;
    private final String methodName;
    private final String className;
    private final boolean isVoidMethod;
    private int paramsVarIndex;
    private int paramsTypeVarIndex;
    private boolean isServiceMethod;
    private int spanIndex;
    private List<Integer> lineList = new ArrayList<>();

    protected MethodInterceptorMethodVisitor(int api, MethodVisitor mv, int access, String name, String descriptor, String className, boolean isServiceMethod) {
        super(api, mv, access, name, descriptor);
        this.methodName = name; // 保存方法名到字段
        this.className = className;
        this.isServiceMethod = isServiceMethod;
        this.isVoidMethod = Type.getReturnType(descriptor).equals(Type.VOID_TYPE);
    }

    @Override
    protected void onMethodEnter() {
        if (isServiceMethod) {
            mv.visitMethodInsn(INVOKESTATIC, Type.getInternalName(ProfilingAspect.class), "init", "()V", false);
            return;
        }
        // 记录方法进入时间
        mv.visitMethodInsn(INVOKESTATIC, SYSTEM_CLASS, CURRENT_TIME_MILLIS_METHOD, "()J", false);
        startTimeVarIndex = newLocal(Type.LONG_TYPE); // 创建新的局部变量来存储时间
        mv.visitVarInsn(LSTORE, startTimeVarIndex);

        // 记录入参
        Type[] argumentTypes = Type.getArgumentTypes(this.methodDesc);
        paramsVarIndex = createArrayForParameters(argumentTypes);
        paramsTypeVarIndex = createArrayTypeForParameters(argumentTypes);

        mv.visitMethodInsn(INVOKESTATIC, Type.getInternalName(ProfilingAspect.class), "before", "()Ljava/lang/String;", false);
        spanIndex = newLocal(Type.getType("Ljava/lang/String;")); // 创建新的局部变量来存储时间
        mv.visitVarInsn(ASTORE, spanIndex);
    }

    @Override
    protected void onMethodExit(int opcode) {
        if ((IRETURN <= opcode && opcode <= RETURN) || opcode == ATHROW) {
            if (isServiceMethod) {
                return;
            }
            // 计算耗时
            mv.visitMethodInsn(INVOKESTATIC, SYSTEM_CLASS, CURRENT_TIME_MILLIS_METHOD, "()J", false);
            mv.visitVarInsn(LLOAD, startTimeVarIndex);
            mv.visitInsn(LSUB);
            int durationVarIndex = newLocal(Type.LONG_TYPE);
            mv.visitVarInsn(LSTORE, durationVarIndex);

            // 准备记录出参
            int returnValueVarIndex = -1;
            Type returnType = Type.getReturnType(this.methodDesc);
            if (!isVoidMethod) {
                switch (opcode) {
                    case DRETURN:
                    case LRETURN:
                        mv.visitInsn(DUP2); // 复制返回值
                        box(returnType); // 装箱基本类型
                        break;
                    case ATHROW:
                        mv.visitInsn(ACONST_NULL); // 异常返回，没有返回值
                        break;
                    case ARETURN:
                    default:
                        mv.visitInsn(DUP); // 复制返回值
                        // 如果返回的是基本类型，则需要进行装箱
                        box(returnType); // 装箱基本类型
                        break;
                }
                returnValueVarIndex = newLocal(Type.getType(Object.class)); // 为返回值分配局部变量
                mv.visitVarInsn(ASTORE, returnValueVarIndex); // 存储返回值/装箱后的对象
            }
            // 调用工具方法打印或记录信息
            mv.visitLdcInsn(className);
            mv.visitLdcInsn(methodName);
            mv.visitVarInsn(LLOAD, durationVarIndex);
            mv.visitLdcInsn(JacksonUtil.toJson(lineList));
            mv.visitVarInsn(ALOAD, paramsTypeVarIndex);
            mv.visitVarInsn(ALOAD, paramsVarIndex);
            mv.visitLdcInsn(returnType.getClassName());
            mv.visitVarInsn(ALOAD, spanIndex);
            if (returnValueVarIndex != -1) {
                mv.visitVarInsn(ALOAD, returnValueVarIndex);
            } else {
                mv.visitInsn(ACONST_NULL);
            }
            mv.visitMethodInsn(INVOKESTATIC, Type.getInternalName(ProfilingAspect.class), "point", "(Ljava/lang/String;Ljava/lang/String;JLjava/lang/String;[Ljava/lang/String;[Ljava/lang/Object;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Object;)V", false);
            mv.visitMethodInsn(INVOKESTATIC, Type.getInternalName(CallContext.class), "exitSpan", "()V", false);
        }
    }

    @Override
    public void visitLineNumber(int line, Label start) {
        super.visitLineNumber(line, start);
        lineList.add(line);
    }

    private int createArrayForParameters(Type[] argumentTypes) {
        // 创建一个新的局部变量来存储参数数组
        int arrayVarIndex = newLocal(Type.getType(Object[].class));
        // 创建一个与参数数量相等长度的数组
        push(argumentTypes.length);
        newArray(Type.getType(Object.class));
        mv.visitVarInsn(ASTORE, arrayVarIndex);

        // 遍历所有参数，将它们存储到数组中
        for (int i = 0; i < argumentTypes.length; i++) {
            mv.visitVarInsn(ALOAD, arrayVarIndex); // 加载数组引用
            push(i); // 加载当前索引
            loadArg(i); // 加载第 i 个参数
            box(argumentTypes[i]); // 如果需要，对基本类型进行装箱
            mv.visitInsn(AASTORE); // 将参数存储到数组中
        }
        return arrayVarIndex;
    }

    private int createArrayTypeForParameters(Type[] argumentTypes) {
        // 创建一个新的局部变量来存储参数数组类型
        int arrayTypeVarIndex = newLocal(Type.getType(String[].class));
        // 创建一个与参数数量相等长度的数组
        push(argumentTypes.length);
        newArray(Type.getType(String.class));
        mv.visitVarInsn(ASTORE, arrayTypeVarIndex);

        // 遍历所有参数，将它们存储到数组中
        for (int i = 0; i < argumentTypes.length; i++) {
            mv.visitVarInsn(ALOAD, arrayTypeVarIndex); // 加载数组引用
            push(i); // 加载当前索引
            mv.visitLdcInsn(argumentTypes[i].getClassName());
            mv.visitInsn(AASTORE); // 将参数存储到数组中
        }
        return arrayTypeVarIndex;
    }
}